import os
import pickle
import sys
import joblib
import numpy as np
import pandas as pd
from flask import Flask, request, render_template, session, redirect, url_for, flash, jsonify, send_from_directory
from functools import wraps
from openai import OpenAI
from werkzeug.utils import secure_filename
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
from collections import OrderedDict

# ==============================================================================
# 0. 初始化 Flask 应用
# ==============================================================================
app = Flask(__name__)
app.config['SECRET_KEY'] = 'sk-a921b8b04f914cc29cee392e32d47a1a'
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# ==============================================================================
# 1. AI & 模型配置
# ==============================================================================
# ... (省略所有模型加载代码，它们都是正确的，无需改动)
# DeepSeek AI 客户端初始化
try:
    DEEPSEEK_API_KEY = 'sk-a921b8b04f914cc29cee392e32d47a1a'
    if '在这里粘贴你自己的' in DEEPSEEK_API_KEY: print("警告：请在 app.py 中设置您的 DeepSeek API 密钥！")
    client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com/v1")
    print("DeepSeek AI 客户端初始化成功！")
except Exception as e:
    print(f"DeepSeek AI 初始化失败: {e}");
    client = None
# 加载标准的 ResNet50 模型并填充权重
try:
    disease_model = models.resnet50()
    num_ftrs = disease_model.fc.in_features
    disease_model.fc = nn.Sequential(nn.Identity(), nn.Linear(num_ftrs, 38))
    state_dict = torch.load('saved_models/model.pth', map_location=torch.device('cpu'))
    disease_model.load_state_dict(state_dict)
    disease_model.eval()
    class_names = ['Apple___Apple_scab', 'Apple___Black_rot', 'Apple___Cedar_apple_rust', 'Apple___healthy',
                   'Blueberry___healthy', 'Cherry_(including_sour)___Powdery_mildew',
                   'Cherry_(including_sour)___healthy', 'Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot',
                   'Corn_(maize)___Common_rust_', 'Corn_(maize)___Northern_Leaf_Blight', 'Corn_(maize)___healthy',
                   'Grape___Black_rot', 'Grape___Esca_(Black_Measles)', 'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)',
                   'Grape___healthy', 'Orange___Haunglongbing_(Citrus_greening)', 'Peach___Bacterial_spot',
                   'Peach___healthy', 'Pepper,_bell___Bacterial_spot', 'Pepper,_bell___healthy',
                   'Potato___Early_blight', 'Potato___Late_blight', 'Potato___healthy', 'Raspberry___healthy',
                   'Soybean___healthy', 'Squash___Powdery_mildew', 'Strawberry___Leaf_scorch', 'Strawberry___healthy',
                   'Tomato___Bacterial_spot', 'Tomato___Early_blight', 'Tomato___Late_blight', 'Tomato___Leaf_Mold',
                   'Tomato___Septoria_leaf_spot', 'Tomato___Spider_mites Two-spotted_spider_mite',
                   'Tomato___Target_Spot', 'Tomato___Tomato_Yellow_Leaf_Curl_Virus', 'Tomato___Tomato_mosaic_virus',
                   'Tomato___healthy']
    translation_dict = {
        'Apple___Apple_scab': '苹果 - 疮痂病', 'Apple___Black_rot': '苹果 - 黑腐病',
        'Apple___Cedar_apple_rust': '苹果 - 雪松锈病', 'Apple___healthy': '苹果 - 健康',
        'Blueberry___healthy': '蓝莓 - 健康', 'Cherry_(including_sour)___Powdery_mildew': '樱桃 - 白粉病',
        'Cherry_(including_sour)___healthy': '樱桃 - 健康',
        'Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot': '玉米 - 灰斑病',
        'Corn_(maize)___Common_rust_': '玉米 - 普通锈病', 'Corn_(maize)___Northern_Leaf_Blight': '玉米 - 北方叶枯病',
        'Corn_(maize)___healthy': '玉米 - 健康', 'Grape___Black_rot': '葡萄 - 黑腐病',
        'Grape___Esca_(Black_Measles)': '葡萄 - 埃斯卡病(黑麻疹)',
        'Grape___Leaf_blight_(Isariopsis_Leaf_Spot)': '葡萄 - 叶枯病', 'Grape___healthy': '葡萄 - 健康',
        'Orange___Haunglongbing_(Citrus_greening)': '橙子 - 黄龙病(柑橘绿化病)',
        'Peach___Bacterial_spot': '桃 - 疮痂病', 'Peach___healthy': '桃 - 健康',
        'Pepper,_bell___Bacterial_spot': '辣椒 - 细菌性斑点病', 'Pepper,_bell___healthy': '辣椒 - 健康',
        'Potato___Early_blight': '马铃薯 - 早疫病', 'Potato___Late_blight': '马铃薯 - 晚疫病',
        'Potato___healthy': '马铃薯 - 健康', 'Raspberry___healthy': '覆盆子 - 健康', 'Soybean___healthy': '大豆 - 健康',
        'Squash___Powdery_mildew': '南瓜 - 白粉病', 'Strawberry___Leaf_scorch': '草莓 - 叶灼病',
        'Strawberry___healthy': '草莓 - 健康', 'Tomato___Bacterial_spot': '番茄 - 细菌性斑点病',
        'Tomato___Early_blight': '番茄 - 早疫病', 'Tomato___Late_blight': '番茄 - 晚疫病',
        'Tomato___Leaf_Mold': '番茄 - 叶霉病', 'Tomato___Septoria_leaf_spot': '番茄 - 斑枯病',
        'Tomato___Spider_mites Two-spotted_spider_mite': '番茄 - 蜘蛛螨(二斑叶螨)',
        'Tomato___Target_Spot': '番茄 - 靶斑病', 'Tomato___Tomato_Yellow_Leaf_Curl_Virus': '番茄 - 黄化曲叶病毒',
        'Tomato___Tomato_mosaic_virus': '番茄 - 花叶病毒', 'Tomato___healthy': '番茄 - 健康'
    }
    preprocess = transforms.Compose([transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])
    print("PyTorch 标准 ResNet50 植物病害识别模型加载成功！")
except Exception as e:
    print(f"PyTorch 模型加载失败: {e}");
    disease_model = None
# 其他模型加载...
import __main__

try:
    from skorch import NeuralNetClassifier
except ImportError:
    print("Warning: skorch not found."); NeuralNetClassifier = None


def to_float32(X): return np.asarray(X, dtype=np.float32)


__main__.to_float32 = to_float32


class _TabTransformer(nn.Module):
    def __init__(self, num_features, num_classes, d_model=64, n_heads=4, num_layers=2, dim_feedforward=128,
                 dropout=0.0):
        super().__init__();
        self.num_features = num_features;
        self.feature_linears = nn.ModuleList([nn.Linear(1, d_model) for _ in range(num_features)]);
        self.feature_emb = nn.Parameter(torch.randn(num_features, d_model) * 0.02);
        self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model));
        nn.init.normal_(self.cls_token, std=0.02);
        enc_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_heads, dim_feedforward=dim_feedforward,
                                               dropout=dropout, batch_first=True, activation="gelu");
        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers);
        self.norm = nn.LayerNorm(d_model);
        self.head = nn.Linear(d_model, num_classes)

    def forward(self, x):
        B = x.size(0);
        tokens = [];
        for i in range(self.num_features): t = self.feature_linears[i](x[:, i:i + 1]);t = t + self.feature_emb[
            i];tokens.append(t.unsqueeze(1))
        x_tokens = torch.cat(tokens, dim=1);
        cls = self.cls_token.repeat(B, 1, 1);
        seq = torch.cat([cls, x_tokens], dim=1);
        enc = self.encoder(seq);
        cls_out = self.norm(enc[:, 0, :]);
        logits = self.head(cls_out);
        return logits


__main__._TabTransformer = _TabTransformer
try:
    with open('saved_models/improved_crop_model.pkl', 'rb') as f:
        crop_recommendation_data = pickle.load(f)
    print("作物推荐模型加载成功!")
except Exception as e:
    print(f"加载作物推荐模型失败: {e}");
    crop_recommendation_data = None


def predict_crop(input_data, model_data):
    if model_data is None: return "模型未加载", ""
    model = model_data['model'];
    scaler = model_data['scaler'];
    label_encoder = model_data['label_encoder'];
    expected_columns = model_data['columns']
    df = pd.DataFrame([input_data], columns=expected_columns).astype(np.float32);
    X_scaled = scaler.transform(df);
    y_pred = model.predict(X_scaled)
    try:
        en_name = label_encoder.inverse_transform(np.asarray(y_pred).astype(int))[0]
    except Exception:
        en_name = str(np.asarray(y_pred)[0])
    cn_map = {'rice': '水稻', 'maize': '玉米', 'chickpea': '鹰嘴豆', 'kidneybeans': '菜豆', 'pigeonpeas': '木豆',
              'mothbeans': '蛾豆', 'mungbean': '绿豆', 'blackgram': '黑豆', 'lentil': '扁豆', 'pomegranate': '石榴',
              'banana': '香蕉', 'mango': '芒果', 'grapes': '葡萄', 'watermelon': '西瓜', 'muskmelon': '甜瓜',
              'apple': '苹果', 'orange': '橙子', 'papaya': '木瓜', 'coconut': '椰子', 'cotton': '棉花', 'jute': '黄麻',
              'coffee': '咖啡'}
    cn_name = cn_map.get(en_name.lower(), f"未知作物 ({en_name})");
    return cn_name, en_name.lower()


from helpers import crop_params, crop_name_map, yield_evaluation


class CropYieldPredictor:
    def __init__(self):
        self.models_dir = "crop_models";
        self.crop_name = None;
        self.crop_chinese_name = None;
        self.model = None;
        self.model_name = "模拟模型";
        self.scaler = None;
        self.has_scaler = False

    def load_model_for_crop(self, crop_name):
        self.crop_name = crop_name;
        self.crop_chinese_name = next((k for k, v in crop_name_map.items() if v == crop_name), "未知")
        if not os.path.exists(self.models_dir): print(
            f"警告: 未找到模型目录 {self.models_dir}，将使用模拟预测。");self.model = None;self.model_name = "模拟模型";return
        model_path_rf = os.path.join(self.models_dir, f"{self.crop_name}_随机森林_model.joblib")
        if os.path.exists(model_path_rf):
            try:
                self.model = joblib.load(model_path_rf);self.model_name = "随机森林";print(
                    f"成功加载 {self.crop_chinese_name} 的随机森林模型。");return
            except Exception as e:
                print(f"加载随机森林模型失败: {e}")
        print(f"未找到 {self.crop_chinese_name} 的真实模型，将使用模拟预测。");
        self.model = None;
        self.model_name = "模拟模型"

    def predict_yield(self, params):
        if self.model is None: return self._simulate_prediction(params)
        try:
            expected_features = ['N', 'P', 'K', 'temperature', 'humidity', 'ph', 'rainfall', 'pesticide']
            input_df = pd.DataFrame([params], columns=expected_features);
            prediction = max(0, self.model.predict(input_df)[0]);
            return {'prediction': prediction, 'min': prediction * 0.95, 'max': prediction * 1.05,
                    'std_dev': prediction * 0.02}
        except Exception as e:
            print(f"真实模型预测失败: {e}，切换到模拟预测。");return self._simulate_prediction(params)

    def _simulate_prediction(self, params: dict):
        import random;
        crop_param = crop_params[self.crop_name];
        base_yield = crop_param['max_yield'] * 0.6
        for param_key, value in params.items():
            if param_key + '_range' in crop_param: param_range = crop_param[param_key + '_range'];optimal = sum(
                param_range) / 2;deviation = abs(value - optimal) / (optimal if optimal != 0 else 1);base_yield *= (
                        1 - deviation * 0.1)
        prediction = max(0, base_yield * random.uniform(0.9, 1.1));
        return {'prediction': prediction, 'min': prediction * 0.95, 'max': prediction * 1.05,
                'std_dev': prediction * 0.02}

    def evaluate_result(self, prediction_result):
        if prediction_result is None: return None
        max_yield = crop_params[self.crop_name]['max_yield'];
        prediction = prediction_result['prediction'];
        percentage = (prediction / max_yield) * 100;
        grade = "很差"
        if percentage >= yield_evaluation['excellent'] * 100:
            grade = "优异"
        elif percentage >= yield_evaluation['good'] * 100:
            grade = "良好"
        elif percentage >= yield_evaluation['average'] * 100:
            grade = "一般"
        elif percentage >= yield_evaluation['poor'] * 100:
            grade = "较差"
        return {'prediction': prediction, 'max_yield': max_yield, 'percentage': percentage, 'grade': grade,
                **prediction_result}


yield_predictor = CropYieldPredictor()


# ==============================================================================
# 2. 核心分析与建议生成函数
# ==============================================================================
def generate_advice_and_scores(user_params, ideal_params):
    advice = []
    scores = []
    param_map = {'N': '氮(N)', 'P': '磷(P)', 'K': '钾(K)', 'temperature': '温度', 'humidity': '湿度', 'ph': 'pH值',
                 'rainfall': '降雨量'}
    if not ideal_params: return [], [], []

    for key, label in param_map.items():
        user_val = user_params.get(key)
        if user_val is None: continue
        ideal_range = ideal_params.get(f"{key}_range")
        if not ideal_range: continue
        min_val, max_val = ideal_range

        if user_val < min_val:
            advice.append(
                f"<li><i class='fa-solid fa-arrow-down'></i> <strong>{label}</strong> 含量偏低 (当前: {user_val}, 建议范围: {min_val}-{max_val})，可能影响生长，建议适当补充。</li>")
        elif user_val > max_val:
            advice.append(
                f"<li><i class='fa-solid fa-arrow-up'></i> <strong>{label}</strong> 含量偏高 (当前: {user_val}, 建议范围: {min_val}-{max_val})，可能造成浪费或抑制其他元素吸收。</li>")
        else:
            advice.append(
                f"<li><i class='fa-solid fa-check'></i> <strong>{label}</strong> 含量非常适宜 (当前: {user_val}, 范围: {min_val}-{max_val})。</li>")

        score = 100
        if user_val < min_val:
            score = max(0, int(100 * (user_val / min_val)))
        elif user_val > max_val:
            score = max(0, int(100 * (max_val / user_val)))
        scores.append(score)
    return advice, scores, list(param_map.values())


# ==============================================================================
# 3. 登录系统逻辑
# ==============================================================================
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'logged_in' not in session: flash('请先登录以访问此页面。', 'warning'); return redirect(url_for('login'))
        return f(*args, **kwargs)

    return decorated_function


@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        if request.form['username'] == 'admin' and request.form['password'] == '123':
            session['logged_in'] = True;
            session['username'] = 'admin';
            flash('登录成功！', 'success');
            return redirect(url_for('home'))
        else:
            flash('用户名或密码错误，请重试。', 'danger')
    return render_template('login.html')


@app.route('/logout')
def logout():
    session.pop('logged_in', None);
    session.pop('username', None);
    flash('您已成功退出登录。', 'info');
    return redirect(url_for('home'))


# ==============================================================================
# 4. Flask 路由定义
# ==============================================================================
@app.route('/')
def home():
    return render_template('home.html', active_page='home')


@app.route('/recommendation')
@login_required
def recommendation_form():
    return render_template('recommend_form.html', active_page='recommendation')


@app.route('/prediction')
@login_required
def prediction_form():
    crop_list = sorted(crop_name_map.keys())
    return render_template('predict_form.html', crop_list=crop_list, active_page='prediction')


@app.route('/disease_detection')
@login_required
def disease_detection_form():
    return render_template('disease_form.html', active_page='suggestion')


@app.route('/ai_chatbot')
@login_required
def ai_chatbot():
    return render_template('ai_assistant.html', active_page='chatbot')


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/detect_disease', methods=['POST'])
@login_required
def detect_disease():
    if 'file' not in request.files: flash('未找到上传文件部分', 'danger'); return redirect(
        url_for('disease_detection_form'))
    file = request.files['file']
    if file.filename == '': flash('未选择任何文件', 'danger'); return redirect(url_for('disease_detection_form'))
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        if not disease_model: flash('病害识别模型未加载，功能不可用。', 'danger'); return redirect(
            url_for('disease_detection_form'))
        try:
            img = Image.open(filepath).convert('RGB')
            input_tensor = preprocess(img)
            input_batch = input_tensor.unsqueeze(0)
            with torch.no_grad():
                output = disease_model(input_batch)
            probabilities = torch.nn.functional.softmax(output, dim=1)
            confidence, pred_idx = torch.max(probabilities, 1)
            en_class_name = class_names[pred_idx.item()]
            cn_class_name = translation_dict.get(en_class_name, en_class_name)
            detection = {"name": cn_class_name, "confidence": f"{confidence.item():.2%}"}
            detections = [detection]
        except Exception as e:
            print(f"图像分类推理失败: {e}");
            flash('图片处理失败，请尝试其他图片。', 'danger');
            return redirect(url_for('disease_detection_form'))
        image_url = url_for('uploaded_file', filename=filename)
        return render_template('result.html', detections=detections, image_file=image_url)
    else:
        flash('只允许上传 png, jpg, jpeg 格式的图片', 'danger')
        return redirect(url_for('disease_detection_form'))


@app.route('/ask_ai', methods=['POST'])
@login_required
def ask_ai():
    if not client: return jsonify({'error': 'DeepSeek AI客户端未初始化，请检查API密钥。'}), 500
    user_question = request.json.get('question')
    if not user_question: return jsonify({'error': '问题不能为空。'}), 400
    try:
        response = client.chat.completions.create(model="deepseek-chat", messages=[{"role": "system",
                                                                                    "content": "你是一位精通农业科学、植物病理学和现代耕作技术的AI农业专家。请用简洁、清晰、专业的语言回答用户的问题。"},
                                                                                   {"role": "user",
                                                                                    "content": user_question}],
                                                  stream=False)
        ai_answer = response.choices[0].message.content
        return jsonify({'answer': ai_answer})
    except Exception as e:
        print(f"调用 DeepSeek API 时出错: {e}");
        return jsonify({'error': 'AI助手当前不可用，请稍后再试。'}), 503


@app.route('/recommend', methods=['POST'])
@login_required
def recommend():
    try:
        user_params = {key: float(request.form[key]) for key in
                       ['N', 'P', 'K', 'temperature', 'humidity', 'ph', 'rainfall']}
        cn_name, en_name = predict_crop(user_params, crop_recommendation_data)

        # [ 修复#2 ] 从 helpers.py 中获取理想参数
        ideal_params = crop_params.get(en_name)

        advice, scores, labels = [], [], []
        if ideal_params:
            advice, scores, labels = generate_advice_and_scores(user_params, ideal_params)

        return render_template('result.html', recommendation=cn_name, advice=advice, chart_labels=labels,
                               chart_scores=scores)
    except Exception as e:
        return f"处理作物推荐时出错: {e}"


@app.route('/predict', methods=['POST'])
@login_required
def predict():
    try:
        crop_cn_name = request.form['crop_type']
        crop_en_name = crop_name_map.get(crop_cn_name)
        if not crop_cn_name: return "无效的作物品种"

        user_params = {key: float(request.form[f'{key}_pred']) for key in
                       ['N', 'P', 'K', 'temperature', 'humidity', 'ph', 'rainfall', 'pesticide']}

        # [ 修复#1 ] 传递正确的英文名给 load_model_for_crop
        yield_predictor.load_model_for_crop(crop_en_name)

        prediction_result = yield_predictor.predict_yield(user_params)
        eval_result = yield_predictor.evaluate_result(prediction_result)
        eval_result['crop_name'] = crop_cn_name
        eval_result['model_used'] = yield_predictor.model_name

        ideal_params = crop_params.get(crop_en_name)
        advice, _, _ = generate_advice_and_scores(user_params, ideal_params)

        return render_template('result.html', prediction=eval_result, advice=advice)
    except Exception as e:
        # 更详细的错误日志
        print(f"处理产量预测时出错: {e}")
        # 返回一个更友好的错误信息
        return f"处理产量预测时出错: {e}"


@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)


# ==============================================================================
# 5. 运行应用
# ==============================================================================
if __name__ == "__main__":
    app.run(debug=True, host='0.0.0.0')